#! /usr/bin/env python3
# vim:fenc=utf-8
import argparse
import os
import subprocess

# mem_stat.py -p linux  -d xxxx

sections_table = {
	"linux": {
		"code": "/^\.text/",
		"const": "/^\.rodata/",
		"Data": "/^\.data/",
		"BSS": "/^\.bss/",
	},
}
def calculate_sizes(files, code_pattern, const_pattern, data_pattern, bss_pattern):
	try:
		# Construct the size command
		size_command = ['size', '-t', '--format=sysv'] + files

		# Construct the awk part of the command
		awk_script = (
			"'{0} {{code_total += $2}}  {1} {{const_total += $2}} {2} {{data_total += $2}} {3} {{bss_total += $2}} END "
			"{{printf \"Total Size of Code: %d bytes\\n\", code_total; "
			"printf \"Total Size of Const: %d bytes\\n\", const_total; "
			"printf \"Total Size of Data: %d bytes\\n\", data_total; "
			"printf \"Total Size of Bss: %d bytes\\n\", bss_total; "
			"printf \"Total Size of Rom: %d bytes\\n\", code_total + const_total ; "
			"printf \"Total Size of RAM: %d bytes\\n\", bss_total + data_total;}}'"
		).format(code_pattern,const_pattern, data_pattern, bss_pattern)
		# Construct the full command by piping size output into awk
		command = size_command + ['|', 'awk', awk_script]
		# print(' '.join(command))

		# Run the command
		result = subprocess.run(' '.join(command), shell=True, capture_output=True, text=True)

		# Display the output
		print(result.stdout)

	except subprocess.CalledProcessError as e:
		print(f"An error occurred: {e}")

def get_sections_table_params(product):
	pattern_list = []


	code_pattern = sections_table[product]['code']
	const_pattern = sections_table[product]['const']
	data_pattern = sections_table[product]['Data']
	bss_pattern = sections_table[product]['BSS']

	pattern_list.append(code_pattern)
	pattern_list.append(const_pattern)
	pattern_list.append(data_pattern)
	pattern_list.append(bss_pattern)

	return pattern_list

def args_parse():
	option = argparse.ArgumentParser(description="Memory Statistics")

	option.add_argument('--verbose',
						action='store_true', help='verbose mode')

	option.add_argument('-p', '--product',
						default='linux', dest="product",
						help='Indicate a product.')


	arg, extra_opts = option.parse_known_args()
	return arg, extra_opts

def main():
	args,extra_opts = args_parse()
	print("\nproduct:{}".format(args.product))
	# print(extra_opts)

	pattern_list = get_sections_table_params(args.product)
	calculate_sizes(extra_opts, pattern_list[0], pattern_list[1], pattern_list[2], pattern_list[3])

if __name__ == "__main__":
	main()

